'''
	Author：Yassin
	用于选取股票价格所属行业，筛选金融行业
	包含函数 Get_background
'''

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt


def Get_background(L: list, Stock_code: list = [], Markettype=[1, 4]) -> pd.DataFrame:
    '''
		查询公司的行业类型
		L:list 为且仅为需要选取的项目 code 的列表，
		Stock_code : list 为所需要选取的股票的代码，默认 = []， 表示查询所有股票
		Markettype : list 为你所需要查询的市场类型，默认为 [1,4], 及沪深 A 股市场
	'''
    df = pd.read_csv("~/Desktop/mpacc2/data_invest/temp/data_lib/back/TRD_Co.csv")

    target_Market = []

    for item in Markettype:
        target_Market.append(df[df["Markettype"] == item])

    df = pd.concat(target_Market).reset_index(drop=True)

    ## 筛选需要的项目
    L = ["Stkcd"] + L
    df = df[L]

    ## 筛选股票代码
    if len(Stock_code) != 0:
        new_df = pd.DataFrame(columns=df.columns)
        for item in Stock_code:
            temp = df[df["Stkcd"] == item]
            new_df = pd.concat([new_df, temp])
        df = new_df.reset_index(drop=True)

    return df

if __name__ == "__main__":
    print(Get_background(["Nnindnme"]))
